from argParser import *
import math
import os
import re
import sys
from dendropy import Tree

ap = ArgParser(sys.argv)

color = ["blue", "red", "green", "purple", "yellow"]


def get_tree_distances(trees):
    tree_distances = []
    for t in trees:
        darr = []
        for t in trees:
            darr.append(None)
        tree_distances.append( darr )
    
    for i in range(0, trees.__len__() ):
        for j in range(0, trees.__len__() ):
            d = trees[i].euclidean_distance( trees[j] )
            tree_distances[i][j] = d
    return tree_distances

def get_dimensions():
    tree_path = ap.getArg("--tree")
    
    fin = open(tree_path, "r")
    lines = fin.readlines()
    fin.close()
    
    trees = []
    
    for l in lines:
        if l.__len__() > 5:
            t = Tree()
            l = l.strip()
            t.read_from_string(l, "newick")
            trees.append(t)
            #print t.__str__()
    get_tree_distances(trees)
    

def print_dimensions_for_trace():
     trace_path = ap.getArg("--trace")
     
     fin = open(trace_path, "r")
     lines = fin.readlines()
     fin.close()
     
     iter_distances = {}
     these_trees = []
     component_count = 0
     tree_count = 0
     found_first = False
     for l in lines:
         if l.__len__() < 2:
             tree_distances = get_tree_distances(these_trees)
             these_trees = []
             tree_count += 1
             #print tree_distances.__str__()
             if component_count == 3:
                 print tree_count,tree_distances[0][2], tree_distances[0][1], tree_distances[1][2]
                 iter_distances[tree_count] = [tree_distances[0][2], tree_distances[0][1], tree_distances[1][2]]
             component_count = 0
         else:   
             l = l.strip()
             if l.startswith("["):
                 l = l.split("]")[1]
             t = Tree()
             t.read_from_string(l, "newick")
             these_trees.append(t)
             component_count += 1
             #print t.__str__() 

     #
     # build x<- lines
     #
     xstr = "x <- c("

     # values
     keys = iter_distances.keys()
     keys.sort()
     for x in keys:
         print "key", x
         this_length = "%.2f"%x
         xstr += this_length + ","
     xstr = re.sub(",$", "", xstr)
     xstr += ")"
     print xstr
     
     valstrs = []
     for i in range(0, iter_distances[1].__len__() ):
         valstrs.append("val" + i.__str__() + " <- c(") 
     for x in keys:
         for i in range(0, iter_distances[1].__len__() ):
             #print "i=",i
             #print "valstrs[i] = ", valstrs[i-1]
             valstrs[i-1] += "%.2f"%iter_distances[x][i] + ","

     for i in range(0, iter_distances[1].__len__() ):
         valstrs[i-1] = re.sub(",$", "", valstrs[i-1])         
     for i in range(0, iter_distances[1].__len__() ):
         valstrs[i-1] += ")"
         print valstrs[i-1].__str__()
         
     fout = open(ap.getArg("--outfile") + ".cran", "w")
     fout.write(xstr + "\n")
     for i in range(0, iter_distances[1].__len__() ):
         fout.write(valstrs[i-1] + "\n")
     fout.write("pdf(\"plot." + ap.getArg("--outfile") + ".pdf\", height = 6, width = 8);\n") 
     fout.write("plot(x,val0, type=\"l\", col=\"" + color[0]+ "\", lwd=3);\n")
     for i in range(1, iter_distances[1].__len__() ):
         fout.write("points(x,val" + i.__str__() + ", type=\"l\", col=\"" + color[i]+ "\", lwd=3);\n")
     fout.write("dev.off()\n")   
     fout.close()
     
     
print_dimensions_for_trace()